#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include "../mbedtls/cipher.h"
#include "../mbedtls/platform.h"

uint8_t key[16] =
    {
        0x06, 0xa9, 0x21, 0x40, 0x36, 0xb8, 0xa1, 0x5b,
        0x51, 0x2e, 0x03, 0xd5, 0x34, 0x12, 0x00, 0x06};

uint8_t iv[16] =
    {
        0x3d, 0xaf, 0xba, 0x42, 0x9d, 0x9e, 0xb4, 0x30,
        0xb4, 0x22, 0xda, 0x80, 0x2c, 0x9f, 0xac, 0x41};

static void dump_buf(char *info, uint8_t *buf, uint32_t len)
{
    mbedtls_printf("%s", info);
    for (int i = 0; i < len; i++)
    {
        // mbedtls_printf("%s%02X%s", i % 16 == 0 ? "\n\t" : " ",
        //                buf[i], i == len - 1 ? "\n" : "");
        mbedtls_printf("%02X", buf[i]);
    }
    mbedtls_printf("\n");
}

int main(void)
{
    int ret = 0;
    mbedtls_cipher_context_t cipher_ctx_enc; /*!<  encryption context      */
    mbedtls_cipher_context_t cipher_ctx_dec; /*!<  decryption context      */

    char *message = "hello world";
    size_t messageLength = strlen(message);

    printf("message len:%zu %s\n", messageLength, message);

    dump_buf("key:", key, sizeof(key));
    dump_buf("iv:", iv, sizeof(iv));

    const mbedtls_cipher_info_t *cipher_info = mbedtls_cipher_info_from_string("AES-128-CBC");

    // entrypt message
    ret = mbedtls_cipher_setup(&cipher_ctx_enc, cipher_info);
    if (ret)
    {
        printf("encrypt cipher setup:%d\n", ret);
        exit(-1);
    }

    ret = mbedtls_cipher_setkey(&cipher_ctx_enc, key, mbedtls_cipher_get_key_bitlen(&cipher_ctx_enc), MBEDTLS_ENCRYPT);
    if (ret)
    {
        printf("encrypt cipher setkey:%d\n", ret);
        exit(-1);
    }

    size_t cipherLen = 0;
    cipherLen += messageLength;
    size_t padding = messageLength % mbedtls_cipher_get_block_size(&cipher_ctx_enc);
    if (padding > 0)
    {
        cipherLen += (mbedtls_cipher_get_block_size(&cipher_ctx_enc) - padding);
    }

    unsigned char *cipherText = (unsigned char *)malloc(cipherLen);
    if (!cipherText)
    {
        printf("cipherText malloc err\n");
        exit(-1);
    }

    cipherLen = 0;
    ret = mbedtls_cipher_crypt(&cipher_ctx_enc,
                               iv,
                               sizeof(iv),
                               message, messageLength,
                               cipherText, &cipherLen);
    if (ret)
    {
        printf("encrypt cipher err:%d\n", ret);
        exit(-1);
    }

    printf("encrypt len:%zu\n", cipherLen);
    dump_buf("cipherText:", cipherText, cipherLen);

    /******************************************************************/
    // decrypt cipherText
    ret = mbedtls_cipher_setup(&cipher_ctx_dec, cipher_info);
    if (ret)
    {
        printf("decrypt cipher setup:%d\n", ret);
        exit(-1);
    }

    ret = mbedtls_cipher_setkey(&cipher_ctx_dec, key, mbedtls_cipher_get_key_bitlen(&cipher_ctx_dec), MBEDTLS_DECRYPT);
    if (ret)
    {
        printf("decrypt cipher setkey:%d\n", ret);
        exit(-1);
    }

    size_t plainLen = 0;
    unsigned char *plainText = (unsigned char *)malloc(cipherLen + sizeof(""));
    if (!plainText)
    {
        printf("plainText malloc err\n");
        exit(-1);
    }

    ret = mbedtls_cipher_crypt(&cipher_ctx_dec,
                               iv,
                               sizeof(iv),
                               cipherText, cipherLen,
                               plainText, &plainLen);
    if (ret)
    {
        printf("encrypt cipher err:%d\n", ret);
        exit(-1);
    }

    // add terminated '\0'
    plainText[plainLen] = '\0';

    printf("decrypt len:%zu %s\n", plainLen, plainText);

    if (cipherText)
    {
        free(cipherText);
    }

    if (plainText)
    {
        free(plainText);
    }

    return ret;
}
